from typing import AsyncGenerator

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from contextlib import asynccontextmanager
import logging

logger = logging.getLogger(__name__)

# 异步数据库 URL 格式：mysql+aiomysql://用户名:密码@地址/数据库名
# 移除time_zone参数，避免与aiomysql版本不兼容
##DATABASE_URL = "mysql+aiomysql://root:Xiongmao12345@10.22.102.58:3306/ai_image?charset=utf8mb4"

DATABASE_URL = "mysql+aiomysql://root:root@127.0.0.1:3306/parse_mp3?charset=utf8mb4"
# 优化连接池配置
async_engine = create_async_engine(
    DATABASE_URL, 
    pool_size=10,  # 减少池大小
    max_overflow=20,  # 减少最大溢出连接
    pool_recycle=3600,  # 1小时回收连接
    pool_timeout=30,  # 30秒超时
    pool_pre_ping=True,  # 添加预ping检查
    echo=False,
    # 添加连接参数，在连接建立后设置时区
    connect_args={
        "init_command": "SET time_zone = '+08:00'",
        "charset": "utf8mb4",
        "autocommit": True
    }
)

# 异步会话工厂
AsyncSessionLocal = sessionmaker(
    bind=async_engine, 
    class_=AsyncSession, 
    expire_on_commit=False
)

Base = declarative_base()

# 获取数据库会话 - 用于FastAPI依赖注入
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    session = AsyncSessionLocal()
    try:
        yield session
        await session.commit()  # 自动提交成功的事务
    except Exception as e:
        await session.rollback()  # 异常时回滚
        logger.error(f"数据库操作异常: {str(e)}")
        raise e
    finally:
        await session.close()  # 确保会话被关闭

# 获取独立数据库会话 - 用于后台任务等
async def get_standalone_db() -> AsyncSession:
    return AsyncSessionLocal()

@asynccontextmanager
async def get_db_session():
    """获取数据库会话的上下文管理器 - 用于service层"""
    session = AsyncSessionLocal()
    try:
        yield session
        await session.commit()
    except Exception as e:
        await session.rollback()
        logger.error(f"数据库操作异常: {str(e)}")
        raise
    finally:
        await session.close()

@asynccontextmanager
async def transaction(session: AsyncSession):
    """带事务管理的上下文管理器"""
    try:
        yield session
        await session.commit()
    except Exception as e:
        await session.rollback()
        raise
    finally:
        await session.close()